"""
游戏AI环境类 - 基于Gymnasium的强化学习环境
"""
import gymnasium as gym
from gymnasium import spaces
import numpy as np
import pygame
from settings import *
import random


class MultiAgentBallBattleEnv(gym.Env):
    """
    多智能体球战游戏环境
    
    动作空间:
    - 0: 不移动，不射击
    - 1: 向上移动
    - 2: 向下移动
    - 3: 向左移动
    - 4: 向右移动
    - 5: 向左上移动
    - 6: 向右上移动
    - 7: 向左下移动
    - 8: 向右下移动
    - 9: 向当前朝向射击
    - 10-19: 向特定方向射击（10个方向）
    
    状态空间:
    - 玩家位置 (x, y)
    - 玩家生命值
    - 所有野怪位置和生命值
    - 所有子弹位置和方向
    - 所有障碍物位置
    """
    
    metadata = {'render_modes': ['human', 'rgb_array'], 'render_fps': FPS}
    
    def __init__(self, render_mode=None):
        super().__init__()
        
        self.render_mode = render_mode
        self.window_size = (WIDTH, HEIGHT)  # 控制窗口大小
        
        # 定义动作空间: 20个离散动作
        # 0-8: 移动动作（包括不移动）
        # 9-19: 射击动作（10个方向）
        self.action_space = spaces.Discrete(20)
        
        # 定义观察空间:
        # - 玩家位置 (2)
        # - 玩家生命值 (1)
        # - 最多10个野怪的位置和生命值 (3 * 10 = 30)
        # - 最多20个子弹的位置和方向 (3 * 20 = 60)
        # - 最多20个障碍物的位置 (2 * 20 = 40)
        # 总计: 2 + 1 + 30 + 60 + 40 = 133
        self.observation_space = spaces.Box(
            low=0, high=max(WIDTH, HEIGHT), 
            shape=(133,), dtype=np.float32
        )
        
        # 初始化游戏状态
        self._action_to_direction = {
            0: np.array([0, 0]),      # 不移动
            1: np.array([0, -1]),     # 向上
            2: np.array([0, 1]),      # 向下
            3: np.array([-1, 0]),     # 向左
            4: np.array([1, 0]),      # 向右
            5: np.array([-1, -1]),    # 向左上
            6: np.array([1, -1]),     # 向右上
            7: np.array([-1, 1]),     # 向左下
            8: np.array([1, 1]),      # 向右下
        }
        
        # 射击方向（10个均匀分布的方向）
        self._action_to_shoot_direction = {}
        for i in range(10):
            angle = i * (360 / 10)  # 每36度一个方向
            rad = np.deg2rad(angle)
            self._action_to_shoot_direction[9 + i] = np.array([np.cos(rad), np.sin(rad)])
        
        # 游戏窗口
        self.window = None
        self.clock = None
        
        # 游戏状态
        self.player_pos = np.array([WIDTH / 2, HEIGHT / 2], dtype=float)
        self.player_health = PLAYER_HEALTH
        self.player_angle = 0  # 玩家朝向
        
        # 野怪列表
        self.monsters = []
        
        # 子弹列表
        self.bullets = []
        
        # 障碍物列表
        self.obstacles = []
        
        # 游戏状态标志
        self.game_over = False
        self.player_won = False
        
        # 计时器
        self.last_shot_time = 0
        self.shot_cooldown = 500  # 射击冷却时间（毫秒）
        
        # 随机数生成器
        self.np_random = None
        self._rng = np.random.default_rng()
    
    def _get_obs(self):
        """获取当前观察状态"""
        obs = np.zeros(133, dtype=np.float32)
        
        # 玩家位置 (0-1)
        obs[0] = self.player_pos[0]
        obs[1] = self.player_pos[1]
        
        # 玩家生命值 (2)
        obs[2] = self.player_health
        
        # 野怪信息 (3-32)
        monster_start = 3
        for i, monster in enumerate(self.monsters[:10]):  # 最多10个野怪
            idx = monster_start + i * 3
            obs[idx] = monster['pos'][0]
            obs[idx + 1] = monster['pos'][1]
            obs[idx + 2] = monster['health']
        
        # 子弹信息 (33-92)
        bullet_start = 33
        for i, bullet in enumerate(self.bullets[:20]):  # 最多20个子弹
            idx = bullet_start + i * 3
            obs[idx] = bullet['pos'][0]
            obs[idx + 1] = bullet['pos'][1]
            # 将方向转换为角度
            angle = np.arctan2(bullet['dir'][1], bullet['dir'][0])
            obs[idx + 2] = angle
        
        # 障碍物信息 (93-132)
        obstacle_start = 93
        for i, obstacle in enumerate(self.obstacles[:20]):  # 最多20个障碍物
            idx = obstacle_start + i * 2
            obs[idx] = obstacle['pos'][0]
            obs[idx + 1] = obstacle['pos'][1]
        
        return obs
    
    def _get_info(self):
        """获取额外信息"""
        return {
            "player_health": self.player_health,
            "monsters_count": len(self.monsters),
            "bullets_count": len(self.bullets),
            "game_over": self.game_over,
            "player_won": self.player_won
        }
    
    def reset(self, seed=None, options=None):
        """重置环境"""
        super().reset(seed=seed)
        
        # 重置随机数生成器
        if seed is not None:
            self.np_random = np.random.default_rng(seed)
        else:
            self.np_random = self._rng
        
        # 重置玩家状态
        self.player_pos = np.array([WIDTH / 2, HEIGHT / 2], dtype=float)
        self.player_health = PLAYER_HEALTH
        self.player_angle = 0
        
        # 重置游戏状态
        self.game_over = False
        self.player_won = False
        self.last_shot_time = 0
        
        # 生成随机野怪
        self.monsters = []
        num_monsters = self.np_random.integers(3, 8)  # 3-7个野怪
        for _ in range(num_monsters):
            # 随机位置，避免太靠近玩家初始位置
            while True:
                x = self.np_random.uniform(50, WIDTH - 50)
                y = self.np_random.uniform(50, HEIGHT - 50)
                # 确保距离玩家初始位置至少200像素
                if np.linalg.norm(np.array([x, y]) - self.player_pos) > 200:
                    break
            
            self.monsters.append({
                'pos': np.array([x, y], dtype=float),
                'health': 50,  # 野怪生命值
                'max_health': 50,
                'speed': MONSTER_SPEED
            })
        
        # 清空子弹
        self.bullets = []
        
        # 生成随机障碍物
        self.obstacles = []
        num_obstacles = self.np_random.integers(5, 10)  # 5-9个障碍物
        for _ in range(num_obstacles):
            # 随机位置和大小
            x = self.np_random.uniform(50, WIDTH - 50)
            y = self.np_random.uniform(50, HEIGHT - 50)
            w = self.np_random.uniform(30, 80)
            h = self.np_random.uniform(30, 80)
            
            # 确保障碍物不与玩家初始位置重叠
            player_rect = pygame.Rect(self.player_pos[0] - PLAYER_SIZE/2, 
                                     self.player_pos[1] - PLAYER_SIZE/2, 
                                     PLAYER_SIZE, PLAYER_SIZE)
            obstacle_rect = pygame.Rect(x - w/2, y - h/2, w, h)
            
            if not player_rect.colliderect(obstacle_rect):
                self.obstacles.append({
                    'pos': np.array([x, y], dtype=float),
                    'width': w,
                    'height': h
                })
        
        observation = self._get_obs()
        info = self._get_info()
        
        if self.render_mode == "human":
            self._render_frame()
        
        return observation, info
    
    def step(self, action):
        """执行一步动作"""
        # 获取当前时间
        current_time = pygame.time.get_ticks()
        
        # 执行动作
        if action < 9:  # 移动动作
            direction = self._action_to_direction[action]
            # 计算新位置
            new_pos = self.player_pos + direction * PLAYER_SPEED / FPS
            
            # 检查边界
            new_pos[0] = np.clip(new_pos[0], PLAYER_SIZE/2, WIDTH - PLAYER_SIZE/2)
            new_pos[1] = np.clip(new_pos[1], PLAYER_SIZE/2, HEIGHT - PLAYER_SIZE/2)
            
            # 检查障碍物碰撞
            player_rect = pygame.Rect(new_pos[0] - PLAYER_SIZE/2, 
                                     new_pos[1] - PLAYER_SIZE/2, 
                                     PLAYER_SIZE, PLAYER_SIZE)
            
            collision = False
            for obstacle in self.obstacles:
                obstacle_rect = pygame.Rect(obstacle['pos'][0] - obstacle['width']/2, 
                                           obstacle['pos'][1] - obstacle['height']/2, 
                                           obstacle['width'], obstacle['height'])
                if player_rect.colliderect(obstacle_rect):
                    collision = True
                    break
            
            if not collision:
                self.player_pos = new_pos
                # 更新玩家朝向
                if action != 0:  # 如果不是不移动
                    self.player_angle = np.arctan2(direction[1], direction[0])
        
        elif 9 <= action < 20:  # 射击动作
            # 检查射击冷却
            if current_time - self.last_shot_time >= self.shot_cooldown:
                shoot_direction = self._action_to_shoot_direction[action]
                
                # 创建新子弹
                self.bullets.append({
                    'pos': self.player_pos.copy(),
                    'dir': shoot_direction,
                    'owner': 'player'
                })
                
                self.last_shot_time = current_time
        
        # 更新子弹位置
        new_bullets = []
        for bullet in self.bullets:
            # 移动子弹
            new_pos = bullet['pos'] + bullet['dir'] * BULLET_SPEED / FPS
            
            # 检查边界
            if (0 <= new_pos[0] <= WIDTH and 0 <= new_pos[1] <= HEIGHT):
                bullet['pos'] = new_pos
                
                # 检查子弹与障碍物碰撞
                bullet_rect = pygame.Rect(new_pos[0] - BULLET_SIZE/2, 
                                        new_pos[1] - BULLET_SIZE/2, 
                                        BULLET_SIZE, BULLET_SIZE)
                
                hit_obstacle = False
                for obstacle in self.obstacles:
                    obstacle_rect = pygame.Rect(obstacle['pos'][0] - obstacle['width']/2, 
                                               obstacle['pos'][1] - obstacle['height']/2, 
                                               obstacle['width'], obstacle['height'])
                    if bullet_rect.colliderect(obstacle_rect):
                        hit_obstacle = True
                        break
                
                if not hit_obstacle:
                    # 检查子弹与野怪碰撞（如果是玩家子弹）
                    if bullet['owner'] == 'player':
                        hit_monster = False
                        for monster in self.monsters:
                            monster_rect = pygame.Rect(monster['pos'][0] - PLAYER_SIZE/2, 
                                                       monster['pos'][1] - PLAYER_SIZE/2, 
                                                       PLAYER_SIZE, PLAYER_SIZE)
                            if bullet_rect.colliderect(monster_rect):
                                # 野怪受伤
                                monster['health'] -= 10
                                hit_monster = True
                                break
                        
                        if not hit_monster:
                            new_bullets.append(bullet)
                    else:
                        new_bullets.append(bullet)
        
        self.bullets = new_bullets
        
        # 更新野怪（简单AI）
        for monster in self.monsters:
            # 简单的追踪玩家逻辑
            direction = self.player_pos - monster['pos']
            distance = np.linalg.norm(direction)
            
            if distance > 0:
                direction = direction / distance  # 归一化
                
                # 移动野怪
                new_pos = monster['pos'] + direction * monster['speed'] / FPS
                
                # 检查边界
                new_pos[0] = np.clip(new_pos[0], PLAYER_SIZE/2, WIDTH - PLAYER_SIZE/2)
                new_pos[1] = np.clip(new_pos[1], PLAYER_SIZE/2, HEIGHT - PLAYER_SIZE/2)
                
                # 检查障碍物碰撞
                monster_rect = pygame.Rect(new_pos[0] - PLAYER_SIZE/2, 
                                          new_pos[1] - PLAYER_SIZE/2, 
                                          PLAYER_SIZE, PLAYER_SIZE)
                
                collision = False
                for obstacle in self.obstacles:
                    obstacle_rect = pygame.Rect(obstacle['pos'][0] - obstacle['width']/2, 
                                               obstacle['pos'][1] - obstacle['height']/2, 
                                               obstacle['width'], obstacle['height'])
                    if monster_rect.colliderect(obstacle_rect):
                        collision = True
                        break
                
                if not collision:
                    monster['pos'] = new_pos
                
                # 野怪攻击玩家（如果足够近）
                if distance < 50:  # 攻击范围
                    if self.np_random.random() < 0.02:  # 2%概率攻击
                        self.player_health -= 5
        
        # 移除死亡的野怪
        self.monsters = [m for m in self.monsters if m['health'] > 0]
        
        # 检查游戏结束条件
        if self.player_health <= 0:
            self.game_over = True
            self.player_won = False
        elif len(self.monsters) == 0:
            self.game_over = True
            self.player_won = True
        
        # 计算奖励
        reward = 0
        
        # 生存奖励
        if not self.game_over:
            reward += 0.1
        
        # 消灭野怪奖励
        if self.game_over and self.player_won:
            reward += 100
        
        # 游戏失败惩罚
        if self.game_over and not self.player_won:
            reward -= 100
        
        observation = self._get_obs()
        info = self._get_info()
        
        terminated = self.game_over
        truncated = False  # 可以添加时间限制等其他截断条件
        
        if self.render_mode == "human":
            self._render_frame()
        
        return observation, reward, terminated, truncated, info
    
    def render(self):
        """渲染环境"""
        if self.render_mode == "rgb_array":
            return self._render_frame()
    
    def _render_frame(self):
        """渲染一帧"""
        if self.window is None and self.render_mode == "human":
            pygame.init()
            pygame.display.init()
            self.window = pygame.display.set_mode(self.window_size)
            pygame.display.set_caption("Multi-Agent Ball Battle - AI Training")
        if self.clock is None and self.render_mode == "human":
            self.clock = pygame.time.Clock()
        
        # 创建画布
        canvas = pygame.Surface(self.window_size)
        canvas.fill(DARKGREY)
        
        # 绘制障碍物
        for obstacle in self.obstacles:
            pygame.draw.rect(canvas, BLUE, 
                           (obstacle['pos'][0] - obstacle['width']/2, 
                            obstacle['pos'][1] - obstacle['height']/2, 
                            obstacle['width'], obstacle['height']))
        
        # 绘制野怪
        for monster in self.monsters:
            # 野怪身体
            pygame.draw.circle(canvas, RED, 
                              (int(monster['pos'][0]), int(monster['pos'][1])), 
                              PLAYER_SIZE // 2)
            
            # 野怪血条
            bar_width = PLAYER_SIZE
            bar_height = 5
            health_percentage = monster['health'] / monster['max_health']
            pygame.draw.rect(canvas, RED, 
                           (monster['pos'][0] - bar_width/2, 
                            monster['pos'][1] - PLAYER_SIZE/2 - 10, 
                            bar_width, bar_height))
            pygame.draw.rect(canvas, GREEN, 
                           (monster['pos'][0] - bar_width/2, 
                            monster['pos'][1] - PLAYER_SIZE/2 - 10, 
                            bar_width * health_percentage, bar_height))
        
        # 绘制玩家
        pygame.draw.circle(canvas, YELLOW, 
                          (int(self.player_pos[0]), int(self.player_pos[1])), 
                          PLAYER_SIZE // 2)
        
        # 玩家血条
        bar_width = PLAYER_SIZE
        bar_height = 5
        health_percentage = self.player_health / PLAYER_HEALTH
        pygame.draw.rect(canvas, RED, 
                       (self.player_pos[0] - bar_width/2, 
                        self.player_pos[1] - PLAYER_SIZE/2 - 10, 
                        bar_width, bar_height))
        pygame.draw.rect(canvas, GREEN, 
                       (self.player_pos[0] - bar_width/2, 
                        self.player_pos[1] - PLAYER_SIZE/2 - 10, 
                        bar_width * health_percentage, bar_height))
        
        # 绘制子弹
        for bullet in self.bullets:
            color = RED if bullet['owner'] == 'player' else WHITE
            pygame.draw.circle(canvas, color, 
                              (int(bullet['pos'][0]), int(bullet['pos'][1])), 
                              BULLET_SIZE // 2)
        
        # 显示游戏状态
        if self.game_over:
            font = pygame.font.Font(None, 36)
            if self.player_won:
                text = font.render("Victory! All monsters eliminated!", True, GREEN)
            else:
                text = font.render("Game Over! Player defeated!", True, RED)
            text_rect = text.get_rect(center=(WIDTH/2, HEIGHT/2))
            canvas.blit(text, text_rect)
        
        if self.render_mode == "human":
            # 将画布复制到窗口
            self.window.blit(canvas, canvas.get_rect())
            pygame.event.pump()
            pygame.display.update()
            self.clock.tick(self.metadata['render_fps'])
        elif self.render_mode == "rgb_array":
            return np.transpose(pygame.surfarray.array3d(canvas), axes=(1, 0, 2))
    
    def close(self):
        """关闭环境"""
        if self.window is not None:
            pygame.display.quit()
            pygame.quit()